
# VPN
import os
os.environ['HTTP_PROXY'] = 'http://127.0.0.1:7890'
os.environ['HTTPS_PROXY'] = 'http://127.0.0.1:7890'

import torch

# 做 分词 和 词典
from transformers import AutoTokenizer
# 词嵌入， 位置编码
from transformers import AutoModel


class TokenizerTool:
    def __init__(self, tokenizer:str ="xlm-roberta-base"):
        self.tokenizer = AutoTokenizer.from_pretrained(tokenizer)

    # 简单的把一个字符串 分割 成一个 str list，里面是 str token
    def split(self, text):
        return self.tokenizer.tokenize(text)

    # 集成了一个 分词器 + 词典
    # 1. 可以输入一个 str , e.g. "hello, world."
    # 2. 也可以输入一个 str list, e.g. ["hello, ...", "coding...", ...]
    def __call__(
            self,
            text: list[str],
            padding: bool = True,
            max_length: int = 1024,
            truncation: bool = True
    ) -> tuple[torch.Tensor, torch.Tensor]:

        res = self.tokenizer(
            text,
            return_tensors="pt",    # 输出 tensor
            padding=padding,        # 开启 padding
            truncation=truncation,  # 默认启用截断（防止过长）
            max_length=max_length   # 可选指定最大长度
        )
        # 返回 int ids 和 padding mask
        return res["input_ids"], res["attention_mask"]


    # 把一个 二维的 tensor int tokens 转为 tokens
    def ids2tokens(self, input_ids):
        if isinstance(input_ids, torch.Tensor):
            input_ids = input_ids.tolist()  # 转成 list[list[int]]
        # input_ids 是 list of list，每一行是一个句子的token id列表
        tokens_batch = []
        for ids in input_ids:
            tokens = self.tokenizer.convert_ids_to_tokens(ids)
            tokens_batch.append(tokens)
        return tokens_batch


    # 把一个 tensor id 转成真正的字符串
    def id2str(self, token_ids):
        return self.tokenizer.batch_decode(token_ids, skip_special_tokens=True)


    # 返回词表 的 大小
    def vocab_size(self) -> int:
        return self.tokenizer.vocab_size


    # 返回 <pad> token 的词表索引
    def pad_id(self) -> int:
        return self.tokenizer.pad_token_id


    # 返回 <bos>（开始）token 的词表索引
    def bos_id(self) -> int:
        return self.tokenizer.bos_token_id


    # 返回 <eos>（结束）token 的词表索引
    def eos_id(self) -> int:
        return self.tokenizer.eos_token_id



if __name__ == "__main__":
    tokenizer = TokenizerTool()

    id, attention_mask = tokenizer(["你好", "hello", "hello, 你好"])
    print(id)
    print(attention_mask)
    print(tokenizer.ids2tokens(id))
    print(tokenizer.id2str(id))

# tensor([[     0,      6, 124084,      2,      1,      1,      1],
#         [     0,  33600,     31,      2,      1,      1,      1],
#         [     0,  33600,     31,      4,      6, 124084,      2]])

# tensor([[1, 1, 1, 1, 0, 0, 0],
#         [1, 1, 1, 1, 0, 0, 0],
#         [1, 1, 1, 1, 1, 1, 1]])

# [['<s>', '▁', '你好', '</s>', '<pad>', '<pad>', '<pad>'],
# ['<s>', '▁hell', 'o', '</s>', '<pad>', '<pad>', '<pad>'],
# ['<s>', '▁hell', 'o', ',', '▁', '你好', '</s>']]


    id, attention_mask = tokenizer(["你好", "hello", "hello, 你好, world。"])
    print(id)
    print(attention_mask)
    print(tokenizer.ids2tokens(id))
    print(tokenizer.id2str(id))

# tensor([[     0,      6, 124084,      2,      1,      1,      1,      1,      1, 1],
#         [     0,  33600,     31,      2,      1,      1,      1,      1,      1, 1],
#         [     0,  33600,     31,      4,      6, 124084,      4,   8999,     30, 2]])

# tensor([[1, 1, 1, 1, 0, 0, 0, 0, 0, 0],
#         [1, 1, 1, 1, 0, 0, 0, 0, 0, 0],
#         [1, 1, 1, 1, 1, 1, 1, 1, 1, 1]])

# [['<s>', '▁', '你好', '</s>', '<pad>', '<pad>', '<pad>', '<pad>', '<pad>', '<pad>'],
# ['<s>', '▁hell', 'o', '</s>', '<pad>', '<pad>', '<pad>', '<pad>', '<pad>', '<pad>'],
# ['<s>', '▁hell', 'o', ',', '▁', '你好', ',', '▁world', '。', '</s>']]

    # 返回 三个 特殊 token 的 id
    print("pad_token_id:", tokenizer.pad_id())
    # 标志一个 token seq 的 起始
    print("bos_token_id:", tokenizer.bos_id())
    # 标志一个 token seq 的 结束
    print("eos_token_id:", tokenizer.eos_id())

    # pad_token_id: 1
    # bos_token_id: 0
    # eos_token_id: 2
